/*
 * Copyright(C) 2022. Huawei Technologies Co.,Ltd. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include <fstream>
#include <iostream>
#include <vector>
#include "BertClassification/BertClassification.h"

namespace {

}

void InitBertParam(InitParam &initParam) {
    initParam.deviceId = 0;
    initParam.labelPath = "./model/bert_text_classification_labels.names";
    initParam.modelPath = "./model/bert_text_classification.om";
    initParam.vocabTextPath = "./data/vocab.txt";
}

int main(int argc, char *argv[])
{
    if (argc <= 1) {
        LogWarn << "Please input text path, such as './BertClassificationV2 ./data/sample.txt'.";
        return APP_ERR_OK;
    }

    InitParam initParam;
    InitBertParam(initParam);
    MxBase::MxInit();
    APP_ERROR ret;

    std::string textPath = argv[1];
    if (textPath.substr(textPath.find_last_of(".")) != ".txt") {
        LogError << "please input the txt file!";
        return APP_ERR_COMM_FAILURE;
    }
    std::string text;
    std::ifstream infile;
    // Open text file.
    infile.open(textPath, std::ios_base::in);
    // Check text file validity.
    if (infile.fail()) {
        LogError << "Failed to open textPath file: " << textPath << ".";
        return APP_ERR_COMM_OPEN_FAIL;
    }

    std::getline(infile, text);
    do {
        if (text.size()==0) {
            LogError << "The text is null, please input right text!";
            continue;
        }
        LogInfo << "origin text: " << text;
        // PreProcess
        std::vector<MxBase::Tensor> inputs;
        ret = BertClassification::TextToTensor(text, initParam, inputs);
        if (ret != APP_ERR_OK) {
            LogError << "BertClassification process failed, ret=" << ret << ".";
            return ret;
        }
        // Inference
        MxBase::Model bert(initParam.modelPath, initParam.deviceId);
        std::vector<MxBase::Tensor> outputs =bert.Infer(inputs);
        for (auto output : outputs) {
            output.ToHost();
        }
        // PostProcess
        std::string label;
        ret = BertClassification::PostProcess(outputs, initParam, label);
        if (ret != APP_ERR_OK) {
            LogError << "BertClassification process failed, ret=" << ret << ".";
            return ret;
        }
        LogInfo << "label: " << label;
    } while (std::getline(infile, text));
    return 0;
}